# -*- encoding: utf-8 -*-
"""
Gini split算法：
1. 分别根据每个特征划分数据集后计算gini split
2. 根据gini split最小值找出最优划分特征
4. 以此类推递归计算每个子数据集并创建递归树，直到分类标签只有一个值
5. 测试递归树
"""


def calculate_gini(dataset):
    """计算数据集的gini系数"""
    gini_prob = 0.0
    # 获取分类标签
    cls_list = [example[-1] for example in dataset]
    cls_set = set(cls_list)
    for c in cls_set:
        gini_prob += cls_list.count(c) / len(cls_list)
    return 1 - gini_prob ** 2


def split_dataset(dataset, index, value):
    """根据特征值划分数据集"""
    sub_dataset = []
    for data in dataset:
        if data[index] == value:
            sub_dataset.append(data[:index] + data[index + 1:])
    return sub_dataset


def find_best_feature(dataset):
    """找出最好的特征，使得根据该特征划分后的gini split达到最小"""
    gini_split = 100
    best_feature_index = -1
    feature_count = len(dataset[0]) - 1
    for i in range(feature_count):
        sub_gini_split = 0.0
        feature_set = set([data[i] for data in dataset])
        for f in feature_set:
            sub_dataset = split_dataset(dataset, i, f)
            gini = calculate_gini(sub_dataset)
            sub_gini_split += len(sub_dataset) / len(dataset) * gini
        if sub_gini_split < gini_split:
            gini_split = sub_gini_split
            best_feature_index = i
    return best_feature_index
